import numpy as np
import pandas as pd
import os
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"  # 必须在导入tensorflow之前！
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_absolute_error
from tensorflow.python.keras.models import Sequential
from tensorflow.python.keras.layers import LSTM, Dense
# from tensorflow.python.keras.optimizers import Adam


# ====== 数据准备 ======
def generate_historical_data(num_items=100):
    """生成模拟历史数据（100个商品90天数据）"""
    np.random.seed(42)
    data = []
    for item_id in range(num_items):
        pre_add = np.random.randint(50, 200)  # 上新前加购数
        daily_add = np.abs(np.random.normal(10, 3, 90))  # 每日加购数
        cum_add = pre_add + np.cumsum(daily_add)

        # 生成销量（前7天较高，后逐渐平稳）
        sales = np.concatenate([
            np.abs(np.random.normal(30, 5, 7)),
            np.abs(np.random.normal(15, 3, 83))
        ])

        for day in range(90):
            data.append([
                item_id, day, daily_add[day], sales[day], cum_add[day]
            ])

    return pd.DataFrame(
        data,
        columns=["item_id", "day", "daily_add", "sales", "cum_add"]
    )


# ====== 特征工程 ======
def create_features(df):
    """构建训练特征"""
    df = df.copy()
    # 滞后特征
    df["last_sales"] = df.groupby("item_id")["sales"].shift(1)
    df["last_add"] = df.groupby("item_id")["daily_add"].shift(1)
    # 转化率特征
    df["conversion_rate"] = df["sales"] / df["cum_add"].replace(0, 1e-5)
    # 填充首日缺失值
    df = df.bfill()
    return df[["day", "last_sales", "last_add", "cum_add"]], df["sales"]


# ====== 模型训练 ======
def train_linear_model(X_train, y_train):
    """训练线性回归模型"""
    model = LinearRegression()
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X_train)
    model.fit(X_scaled, y_train)
    return model, scaler


def train_lstm_model(X_train, y_train, window_size=7):
    """训练LSTM模型（适合长期预测）"""
    # 数据重组为3D格式 [样本数, 时间步, 特征数]
    X_lstm = []
    y_lstm = []
    for i in range(window_size, len(X_train)):
        X_lstm.append(X_train[i - window_size:i])
        y_lstm.append(y_train.iloc[i])

    X_lstm = np.array(X_lstm)
    y_lstm = np.array(y_lstm)

    # 构建LSTM网络
    model = Sequential([
        LSTM(64, input_shape=(window_size, X_train.shape[1]), return_sequences=True),
        LSTM(32),
        Dense(1)
    ])
    # optimizer =
    model.compile(loss="mse")
    # model.compile( Adam(0.001),loss="mse")
    model.fit(X_lstm, y_lstm, epochs=50, batch_size=32, verbose=0)
    return model


# ====== 滚动预测 ======
def rolling_forecast(new_item_data, historical_model, scaler):
    """执行14天滚动预测（每日更新模型）"""
    predictions = []
    actuals = []
    model = historical_model

    for day in range(1, 15):
        # 准备当日特征
        current_day = day - 1
        features = new_item_data[new_item_data["day"] == current_day]
        X, _ = create_features(features)
        X_scaled = scaler.transform(X)

        # 预测次日销量
        pred = model.predict(X_scaled)[0]
        predictions.append(pred)

        # 模拟获取当日真实数据（实际应用中替换为真实API）
        actual = get_real_sales(new_item_data, day)  # 模拟函数
        actuals.append(actual)

        # 更新数据集
        new_row = pd.DataFrame([{
            "item_id": 999,  # 新品ID
            "day": day,
            "daily_add": np.random.randint(5, 20),  # 模拟当日加购
            "sales": actual,
            "cum_add": new_item_data["cum_add"].iloc[-1] + new_item_data["daily_add"].iloc[-1]
        }])
        new_item_data = pd.concat([new_item_data, new_row])

        # 每日更新模型（增量训练）
        if day % 3 == 0:  # 每3天全量更新一次
            X_all, y_all = create_features(new_item_data)
            X_scaled_all = scaler.transform(X_all)
            model.fit(X_scaled_all, y_all)

    return predictions, actuals


# ====== 模拟真实数据获取 ======
def get_real_sales(data, day):
    """模拟真实销量生成（实际应用替换为API调用）"""
    base_sales = 20 if day < 7 else 15  # 首周销量较高
    return max(0, np.random.normal(base_sales, 3))


# ====== 主执行流程 ======
if __name__ == "__main__":
    # 1. 准备数据
    historical_data = generate_historical_data()
    X_hist, y_hist = create_features(historical_data)

    # 2. 训练基础模型
    linear_model, scaler = train_linear_model(X_hist, y_hist)
    lstm_model = train_lstm_model(X_hist, y_hist)

    # 3. 初始化新品数据
    new_item_init = pd.DataFrame([{
        "item_id": 999,
        "day": 0,
        "daily_add": 35,  # 上市日加购
        "sales": 28,  # 上市日销量
        "cum_add": 120  # 上新前累计加购+当日加购
    }])

    # 4. 执行滚动预测
    linear_preds, actuals = rolling_forecast(
        new_item_init.copy(), linear_model, scaler
    )

    # 5. 结果评估
    mae = mean_absolute_error(actuals, linear_preds)
    print(f"预测结果（14天MAE={mae:.2f}）:")
    print(f"实际销量: {actuals}")
    print(f"预测销量: {[round(p, 1) for p in linear_preds]}")

    # 6. 可视化（可选）
    import matplotlib.pyplot as plt

    plt.figure(figsize=(10, 4))
    plt.plot(actuals, label="Actual Sales", marker="o")
    plt.plot(linear_preds, label="Predicted Sales", linestyle="--")
    plt.title("14-Day Rolling Sales Forecast")
    plt.xlabel("Days After Launch")
    plt.ylabel("Sales Quantity")
    plt.legend()
    plt.grid(True)
    plt.show()